package com.yupi.aicodehelper.ai.dto;

import dev.langchain4j.model.chat.StreamingChatModel;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.model.ModelProvider;
import dev.langchain4j.model.chat.listener.ChatModelListener;
import dev.langchain4j.model.chat.request.ChatRequest;
import dev.langchain4j.model.chat.request.ChatRequestParameters;
import dev.langchain4j.model.chat.request.DefaultChatRequestParameters;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.CompleteToolCall;
import dev.langchain4j.model.chat.response.PartialThinking;
import dev.langchain4j.model.chat.response.PartialToolCall;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import java.util.List;

public class MonitoringStreamingChatModel implements StreamingChatModel {

    private final StreamingChatModel delegate;
    private final long firstTokenThresholdMs;   // 首 token 阈值
    private final long avgTokenThresholdMs;     // 平均 token 阈值

    public MonitoringStreamingChatModel(StreamingChatModel delegate,
                                        long firstTokenThresholdMs,
                                        long avgTokenThresholdMs) {
        this.delegate = delegate;
        this.firstTokenThresholdMs = firstTokenThresholdMs;
        this.avgTokenThresholdMs = avgTokenThresholdMs;
    }

    @Override
    public void chat(ChatRequest request, StreamingChatResponseHandler handler) {
        long start = System.currentTimeMillis();
        final long[] firstTokenTime = {0};
        final int[] tokenCount = {0};

        delegate.chat(request, new StreamingChatResponseHandler() {
            @Override
            public void onPartialResponse(String partialResponse) {
                tokenCount[0]++;
                if (firstTokenTime[0] == 0) {
                    firstTokenTime[0] = System.currentTimeMillis();
                    long firstLatency = firstTokenTime[0] - start;
                    if (firstLatency > firstTokenThresholdMs) {
                        System.err.println("[WARN] 首 token 延迟过高: " + firstLatency + " ms");
                    } else {
                        System.out.println("[INFO] 首 token 延迟: " + firstLatency + " ms");
                    }
                }
                handler.onPartialResponse(partialResponse);
            }

            @Override
            public void onCompleteResponse(ChatResponse completeResponse) {
                long end = System.currentTimeMillis();
                if (tokenCount[0] > 1) {
                    long avg = (end - firstTokenTime[0]) / (tokenCount[0] - 1);
                    if (avg > avgTokenThresholdMs) {
                        System.err.println("[WARN] 平均 token 时间过高: " + avg + " ms (" + tokenCount[0] + " tokens)");
                    } else {
                        System.out.println("[INFO] 平均 token 时间: " + avg + " ms (" + tokenCount[0] + " tokens)");
                    }
                }

                handler.onCompleteResponse(completeResponse);
            }

            @Override
            public void onError(Throwable error) {
                handler.onError(error);
            }
        });
    }
}


